package com.initcap.saml;
/**
 * @author initcap
 * @date Created in 2024/6/14 20:36.
 */

import java.util.HashSet;
import java.util.List;
import java.util.Set;
import org.springframework.context.annotation.Configuration;
import org.springframework.core.convert.converter.Converter;
import org.springframework.security.authentication.ProviderManager;
import org.springframework.security.config.annotation.web.builders.HttpSecurity;
import org.springframework.security.config.annotation.web.configuration.WebSecurityConfigurerAdapter;
import org.springframework.security.core.GrantedAuthority;
import org.springframework.security.core.authority.SimpleGrantedAuthority;
import org.springframework.security.saml2.provider.service.authentication.OpenSamlAuthenticationProvider;
import org.springframework.security.saml2.provider.service.authentication.Saml2AuthenticatedPrincipal;
import org.springframework.security.saml2.provider.service.authentication.Saml2Authentication;

@Configuration
public class SecurityConfiguration extends WebSecurityConfigurerAdapter {

    @Override
    protected void configure(HttpSecurity httpSecurity) throws Exception {
        // This class is deprecated, but you have to use it if you're using OpenSAML < 4.0
        OpenSamlAuthenticationProvider authenticationProvider = new OpenSamlAuthenticationProvider();
        authenticationProvider.setAssertionValidator(OpenSamlAuthenticationProvider.createDefaultAssertionValidator());
        authenticationProvider.setResponseAuthenticationConverter(OpenSamlAuthenticationProvider.createDefaultResponseAuthenticationConverter());
        authenticationProvider.setResponseAuthenticationConverter(groupsConverter());

        httpSecurity.authorizeRequests(authz -> authz.anyRequest().authenticated())
                .saml2Login(saml2 -> saml2.authenticationManager(new ProviderManager(authenticationProvider)));
    }

    private Converter<OpenSamlAuthenticationProvider.ResponseToken, Saml2Authentication> groupsConverter() {

        Converter<OpenSamlAuthenticationProvider.ResponseToken, Saml2Authentication> delegate =
                OpenSamlAuthenticationProvider.createDefaultResponseAuthenticationConverter();

        return (responseToken) -> {
            Saml2Authentication authentication = delegate.convert(responseToken);
            Saml2AuthenticatedPrincipal principal = (Saml2AuthenticatedPrincipal) authentication.getPrincipal();
            List<String> groups = principal.getAttribute("groups");
            Set<GrantedAuthority> authorities = new HashSet<>();
            if (groups != null) {
                groups.stream().map(SimpleGrantedAuthority::new).forEach(authorities::add);
            } else {
                authorities.addAll(authentication.getAuthorities());
            }
            return new Saml2Authentication(principal, authentication.getSaml2Response(), authorities);
        };
    }
}
